import tkinter as tk
from tkinter import ttk, filedialog
import torch
import numpy as np
import json
from datetime import datetime
from threading import Thread
import os
from PIL import Image, ImageDraw, ImageTk
import torch.cuda.amp as amp
import math
from collections import defaultdict

# 添加SE注意力模块
class SELayer(torch.nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.avg_pool = torch.nn.AdaptiveAvgPool2d(1)
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(channel, channel // reduction),
            torch.nn.LeakyReLU(0.1),
            torch.nn.Linear(channel // reduction, channel),
            torch.nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

# 修改GomokuNet类与训练代码一致
class GomokuNet(torch.nn.Module):
    def __init__(self, res_blocks=10):  # 参数与训练代码一致
        super().__init__()
        self.conv_init = torch.nn.Conv2d(2, 256, 3, padding=1)  # 通道数256
        self.res_blocks = torch.nn.Sequential(*[
            torch.nn.Sequential(
                torch.nn.Conv2d(256, 256, 3, padding=1),
                torch.nn.BatchNorm2d(256),
                torch.nn.LeakyReLU(0.1),
                SELayer(256),  # 添加SE模块
                torch.nn.Conv2d(256, 256, 3, padding=1),
                torch.nn.BatchNorm2d(256),
                SELayer(256)
            ) for _ in range(res_blocks)
        ])
        # 策略头与训练代码一致
        self.policy_head = torch.nn.Sequential(
            torch.nn.Conv2d(256, 64, 1),
            torch.nn.LeakyReLU(0.1),
            torch.nn.Dropout(0.2),
            torch.nn.Flatten(),
            torch.nn.Linear(64*15*15, 512),
            torch.nn.LeakyReLU(0.1),
            torch.nn.Linear(512, 225)
        )
        # 价值头与训练代码一致
        self.value_head = torch.nn.Sequential(
            torch.nn.Conv2d(256, 32, 1),
            torch.nn.LeakyReLU(0.1),
            torch.nn.AdaptiveAvgPool2d(1),
            torch.nn.Flatten(),
            torch.nn.Linear(32, 64),
            torch.nn.LeakyReLU(0.1),
            torch.nn.Linear(64, 1)
        )

    def forward(self, x):
        x = torch.relu(self.conv_init(x))
        residual = x
        for block in self.res_blocks:
            x = torch.relu(block(x) + residual)
            residual = x
        return (
            torch.log_softmax(self.policy_head(x), dim=1),
            torch.tanh(self.value_head(x))
        )

# MCTS相关类
class MCTSNode:
    def __init__(self, parent=None, prior=1.0):
        self.parent = parent
        self.children = {}
        self.visit_count = 0
        self.total_value = 0.0
        self.prior = prior
        self.state = None

    def expanded(self):
        return len(self.children) > 0

    def value(self):
        return self.total_value / self.visit_count if self.visit_count else 0

class MCTS:
    def __init__(self, model, device, simulations=800, c_puct=4.0):
        self.model = model
        self.device = device
        self.simulations = simulations
        self.c_puct = c_puct
        self.root = MCTSNode()

    def search(self, initial_state):
        for _ in range(self.simulations):
            node = self.root
            state = initial_state.copy()

            # 选择
            while node.expanded():
                action, node = self.select_child(node)
                state = self.apply_action(state, action)

            # 扩展
            if self.is_terminal(state):
                value = self.get_reward(state)
            else:
                self.expand_node(node, state)
                action, node = self.select_child(node)
                state = self.apply_action(state, action)
                value = self.simulate(state)

            # 反向传播
            self.backpropagate(node, value)

    def select_child(self, node):
        _, action, child = max(
            (self.ucb_score(node, child), action, child)
            for action, child in node.children.items()
        )
        return action, child

    def ucb_score(self, parent, child):
        prior_score = child.prior * math.sqrt(parent.visit_count) / (1 + child.visit_count)
        value_score = child.value()
        return value_score + self.c_puct * prior_score

    def expand_node(self, node, state):
        with torch.no_grad(), torch.amp.autocast('cuda'):
            current = (state == 1).astype(np.float32)
            opponent = (state == 2).astype(np.float32)
            state_tensor = torch.tensor(
                np.stack([current, opponent]),
                dtype=torch.float32,
                device=self.device
            ).unsqueeze(0)

            policy_logits, value = self.model(state_tensor)
            policy_probs = torch.softmax(policy_logits[0].cpu(), dim=0).numpy()

        valid_moves = (state.reshape(-1) == 0)
        for action in np.nonzero(valid_moves)[0]:
            child = MCTSNode(parent=node, prior=policy_probs[action])
            child.state = self.apply_action(state, action)
            node.children[action] = child

    def simulate(self, state):
        with torch.no_grad(), torch.amp.autocast('cuda'):
            current = (state == 1).astype(np.float32)
            opponent = (state == 2).astype(np.float32)
            state_tensor = torch.tensor(
                np.stack([current, opponent]),
                dtype=torch.float32,
                device=self.device
            ).unsqueeze(0)

            _, value = self.model(state_tensor)
            return value.item()

    def backpropagate(self, node, value):
        while node is not None:
            node.visit_count += 1
            node.total_value += value
            node = node.parent

    def apply_action(self, state, action):
        new_state = state.copy()
        x, y = divmod(action, 15)
        new_state[y][x] = 2 if state.sum() % 2 == 0 else 1
        return new_state

    def is_terminal(self, state):
        for x in range(15):
            for y in range(15):
                if state[y][x] != 0:
                    if self._check_winner(x, y, state[y][x], state):
                        return True
        return np.all(state != 0)

    def get_reward(self, state):
        for x in range(15):
            for y in range(15):
                if state[y][x] != 0:
                    if self._check_winner(x, y, state[y][x], state):
                        return 1 if state[y][x] == 2 else -1
        return 0

    def _check_winner(self, x, y, player, state):
        directions = [(1, 0), (0, 1), (1, 1), (1, -1)]
        for dx, dy in directions:
            count = 1
            i, j = x + dx, y + dy
            while 0 <= i < 15 and 0 <= j < 15 and state[j][i] == player:
                count += 1
                i += dx
                j += dy
            i, j = x - dx, y - dy
            while 0 <= i < 15 and 0 <= j < 15 and state[j][i] == player:
                count += 1
                i -= dx
                j -= dy
            if count >= 5:
                return True
        return False

    def get_move_probs(self, temperature=1.0):
        visit_counts = [self.root.children[a].visit_count if a in self.root.children else 0
                        for a in range(225)]
        probs = np.power(visit_counts, 1.0 / temperature)
        probs /= probs.sum()
        return probs

# 修改GomokuAI类的加载方法
class GomokuAI:
    def __init__(self):
        self.model = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.last_eval = 0.0
        self.mcts = None

    def load_model(self, model_path):
        try:
            checkpoint = torch.load(model_path, map_location=self.device)
            # 从checkpoint获取配置参数
            config = checkpoint.get('config', {'res_blocks': 10})  # 默认值改为10
            # 动态创建模型
            self.model = GomokuNet(res_blocks=config['res_blocks']).to(self.device)
            # 兼容新旧模型保存格式
            if 'model_state' in checkpoint:
                self.model.load_state_dict(checkpoint['model_state'])
            else:  # 兼容旧版直接保存state_dict的情况
                self.model.load_state_dict(checkpoint)
            self.model.eval()
            self.mcts = MCTS(self.model, self.device)
            return True
        except Exception as e:
            print(f"加载模型失败: {str(e)}")
            return False

    def get_move(self, board):
        if not self.model or not self.mcts:
            return None, 0.0

        # 使用MCTS进行搜索
        self.mcts.root = MCTSNode()  # 重置搜索树
        self.mcts.root.state = board.copy()
        self.mcts.search(board)

        # 获取MCTS建议的走子概率
        move_probs = self.mcts.get_move_probs(temperature=0.1)

        valid_moves = (board.reshape(-1) == 0)
        policy_probs = move_probs * valid_moves

        if policy_probs.sum() > 0:
            policy_probs /= policy_probs.sum()
        else:
            # 处理全部无效的情况
            policy_probs = valid_moves / valid_moves.sum()

        # 选择概率最高的有效落子点
        action = np.argmax(policy_probs)
        x, y = divmod(action, 15)  # 注意坐标顺序调整

        return (x, y), self.mcts.root.value()

class GomokuUI:
    def __init__(self):
        self.window = tk.Tk()
        self.window.title("五子棋 AI Pro")
        self.window.geometry("900x680")
        self.stone_images = {}
        self.move_history = []
        self.game_active = True
        self._init_ui()
        self.new_game()
        self._update_eval()
        self.window.mainloop()

    def _init_ui(self):
        self.cell_size = 40
        self.canvas = tk.Canvas(self.window,
                                width=15 * self.cell_size,
                                height=15 * self.cell_size,
                                bg="#DEB887")
        self.canvas.pack(side=tk.LEFT, padx=20, pady=20)
        self._draw_board()

        control_frame = ttk.Frame(self.window, padding=20)
        control_frame.pack(side=tk.RIGHT, fill=tk.Y)

        ttk.Button(control_frame, text="加载AI模型", command=self._load_model).pack(pady=10)
        self.model_label = ttk.Label(control_frame, text="当前模型: 无")
        self.model_label.pack()

        self.status_label = ttk.Label(control_frame,
                                      text="玩家回合",
                                      font=('微软雅黑', 14))
        self.status_label.pack(pady=20)

        self.eval_label = ttk.Label(control_frame,
                                    text="局面评估: 0.0%",
                                    font=('微软雅黑', 12))
        self.eval_label.pack()

        ttk.Button(control_frame, text="新游戏", command=self.new_game).pack(pady=10)
        ttk.Button(control_frame, text="保存棋谱", command=self._save_game).pack(pady=5)
        ttk.Button(control_frame, text="加载棋谱", command=self._load_game).pack(pady=5)

        self.canvas.bind("<Motion>", self._on_hover)
        self.canvas.bind("<Leave>", self._clear_hover)
        self.canvas.bind("<Button-1>", self._on_click)

    def _draw_board(self):
        offset = self.cell_size // 2
        for i in range(15):
            self.canvas.create_line(offset, offset + i * self.cell_size,
                                    offset + 14 * self.cell_size, offset + i * self.cell_size,
                                    width=1.5, fill="#5E3B23")
            self.canvas.create_line(offset + i * self.cell_size, offset,
                                    offset + i * self.cell_size, offset + 14 * self.cell_size,
                                    width=1.5, fill="#5E3B23")
        stars = [(3, 3), (3, 11), (11, 3), (11, 11), (7, 7)]
        for x, y in stars:
            cx = offset + x * self.cell_size
            cy = offset + y * self.cell_size
            self.canvas.create_oval(cx - 4, cy - 4, cx + 4, cy + 4, fill="#5E3B23", outline="")

    def new_game(self):
        self.board = np.zeros((15, 15), dtype=int)
        self.current_player = 1
        self.game_active = True
        self.move_history = []
        self.canvas.delete("stone")
        self.canvas.delete("win_line")
        self.status_label.config(text="玩家回合")
        self.eval_label.config(text="局面评估: 0.0%")

    def _load_model(self):
        path = filedialog.askopenfilename(filetypes=[("PyTorch Model", "*.pth")])
        if path:
            self.ai = GomokuAI()
            if self.ai.load_model(path):
                self.model_label.config(text=f"当前模型: {os.path.basename(path)}")

    def _update_eval(self):
        if hasattr(self, 'ai'):
            eval_percent = (self.ai.last_eval + 1) * 50
            self.eval_label.config(text=f"局面评估: {eval_percent:.1f}%")
        self.window.after(500, self._update_eval)

    def _create_stone_image(self, color):
        size = self.cell_size - 4
        img = Image.new("RGBA", (size, size))
        draw = ImageDraw.Draw(img)

        if color == "black":
            base = (45, 45, 45)
            highlight = (100, 100, 100)
        else:
            base = (245, 235, 220)
            highlight = (255, 255, 255)

        draw.ellipse((0, 0, size, size), fill=base)
        max_i = int(size * 0.75)
        for i in range(size // 2, max_i):
            progress = (i - size // 2) / (max_i - size // 2)
            alpha = int(120 * (1 - progress))
            start = int(size // 3 * (1 - progress))
            end = size - start
            draw.ellipse(
                (max(0, start), max(0, start),
                 min(size, end), min(size, end)),
                fill=highlight + (alpha,),
                width=2
            )
        center_size = size // 4
        draw.ellipse(
            (center_size, center_size,
             size - center_size, size - center_size),
            fill=highlight + (80,),
            outline=highlight
        )
        return ImageTk.PhotoImage(img)

    def _draw_preview_stone(self, x, y):
        offset = self.cell_size // 2
        cx = offset + x * self.cell_size
        cy = offset + y * self.cell_size

        if "preview" not in self.stone_images:
            img = Image.new("RGBA", (self.cell_size - 4, self.cell_size - 4))
            draw = ImageDraw.Draw(img)
            draw.ellipse((0, 0, self.cell_size - 4, self.cell_size - 4), fill=(0, 0, 0, 128))
            self.stone_images["preview"] = ImageTk.PhotoImage(img)

        self.canvas.delete("preview")
        self.canvas.create_image(cx, cy,
                                 image=self.stone_images["preview"],
                                 anchor=tk.CENTER, tags="preview")

    def _place_stone(self, x, y, player):
        offset = self.cell_size // 2
        cx = offset + x * self.cell_size
        cy = offset + y * self.cell_size

        color = "black" if player == 1 else "white"
        if color not in self.stone_images:
            self.stone_images[color] = self._create_stone_image(color)

        self.canvas.create_image(cx, cy,
                                 image=self.stone_images[color],
                                 anchor=tk.CENTER, tags="stone")
        self.board[y][x] = player
        self.move_history.append((player, x, y))

        if self._check_winner(x, y, player):
            self.game_active = False
            self.status_label.config(text=f"玩家{player} 获胜！")
            self._show_win_line(x, y, player)

    def _check_winner(self, x, y, player):
        directions = [(1, 0), (0, 1), (1, 1), (1, -1)]
        for dx, dy in directions:
            count = 1
            i, j = x + dx, y + dy
            while 0 <= i < 15 and 0 <= j < 15 and self.board[j][i] == player:
                count += 1
                i += dx
                j += dy
            i, j = x - dx, y - dy
            while 0 <= i < 15 and 0 <= j < 15 and self.board[j][i] == player:
                count += 1
                i -= dx
                j -= dy
            if count >= 5:
                return True
        return False

    def _show_win_line(self, x, y, player):
        offset = self.cell_size // 2
        cx = offset + x * self.cell_size
        cy = offset + y * self.cell_size
        self.canvas.create_line(cx - 15, cy, cx + 15, cy,
                                fill="red", width=3, tags="win_line")
        self.canvas.create_line(cx, cy - 15, cx, cy + 15,
                                fill="red", width=3, tags="win_line")

    def _save_game(self):
        path = filedialog.asksaveasfilename(
            defaultextension=".json",
            filetypes=[("棋谱文件", "*.json")]
        )
        if path:
            game_data = {
                "date": datetime.now().isoformat(),
                "moves": self.move_history,
                "board_size": 15
            }
            with open(path, 'w') as f:
                json.dump(game_data, f, indent=2)

    def _load_game(self):
        path = filedialog.askopenfilename(filetypes=[("棋谱文件", "*.json")])
        if path:
            try:
                with open(path) as f:
                    data = json.load(f)
                self.new_game()
                for player, x, y in data["moves"]:
                    self._place_stone(x, y, player)
                    self.current_player = 2 if player == 1 else 1
            except Exception as e:
                print(f"加载棋谱失败: {str(e)}")

    def _on_hover(self, event):
        x = (event.x - self.cell_size // 2) // self.cell_size
        y = (event.y - self.cell_size // 2) // self.cell_size
        if 0 <= x < 15 and 0 <= y < 15 and self.board[y][x] == 0 and self.game_active:
            self._draw_preview_stone(x, y)

    def _clear_hover(self, event):
        self.canvas.delete("preview")

    def _on_click(self, event):
        if self.current_player == 1 and self.game_active:
            x = (event.x - self.cell_size // 2) // self.cell_size
            y = (event.y - self.cell_size // 2) // self.cell_size
            if 0 <= x < 15 and 0 <= y < 15 and self.board[y][x] == 0:
                self._place_stone(x, y, 1)
                self.current_player = 2
                self.status_label.config(text="AI思考中...")
                Thread(target=self._ai_move).start()

    def _ai_move(self):
        if hasattr(self, 'ai') and self.game_active:
            move, _ = self.ai.get_move(self.board)
            if move:
                x, y = move
                self._place_stone(x, y, 2)
                self.current_player = 1
                self.status_label.config(text="玩家回合")

if __name__ == "__main__":
    # print("欢迎使用五子棋游戏！")
    GomokuUI()